import torch
import torch.nn as nn
import torch.optim as optim
import torchtext
import spacy
import random
from torchtext.data.utils import get_tokenizer
from torchtext.vocab import build_vocab_from_iterator
from torchtext.transforms import VocabTransform
import torch
from torchtext.datasets import IMDB
from torch.utils.data import DataLoader
# 定义数据处理函数


def yield_tokens(data_iter):
    for _, text in data_iter:
        yield tokenizer(text)


# 定义分词器
tokenizer = get_tokenizer("spacy", language="en_core_web_sm")
# 加载 IMDB 数据集
train_iter, test_iter = IMDB(split=("train", "test"))

# 加载 spaCy 的英文模型
nlp = spacy.load('en_core_web_sm')
# 构建词汇表
vocab = build_vocab_from_iterator(yield_tokens(train_iter), specials=["<unk>", "<pad>", "<bos>", "<eos>"])
vocab.set_default_index(vocab["<unk>"])  # 设置默认索引为 <unk>

# 定义文本转换函数
text_transform = VocabTransform(vocab)

# 定义分词函数
def tokenize(text):
    return [token.text for token in nlp(text.lower())]


# 对数据集中的文本进行分词
def preprocess_dataset(example):
    example['text'] = tokenize(example['text'])
    return example


# 应用分词函数到数据集
dataset = dataset.map(preprocess_dataset)
print(dataset)
exit(0)

SEED = 1234
MAX_VOCAB_SIZE = 25000
BATCH_SIZE = 64
torch.manual_seed(SEED)
torch.backends.cudnn.deterministic = True

TEXT = torchtext.data.Field(tokenize='spacy', tokenizer_language='en_core_web_sm', lower=True)
LABEL = torchtext.data.LabelField(dtype=torch.float)

train_data, test_data = torchtext.datasets.IMDB.splits(TEXT, LABEL)

print(f'Number of training examples: {len(train_data)}')
print(f'Number of testing examples: {len(test_data)}')

TEXT.build_vocab(train_data, max_size=MAX_VOCAB_SIZE)
LABEL.build_vocab(train_data)

print(f"Unique tokens in TEXT vocabulary: {len(TEXT.vocab)}")
print(f"Unique tokens in LABEL vocabulary: {len(LABEL.vocab)}")

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

train_iterator, test_iterator = torchtext.data.BucketIterator.splits(
    (train_data, test_data),
    batch_size=BATCH_SIZE,
    device=device)


class TransformerModel(nn.Module):
    def __init__(self, input_dim, emb_dim, n_heads, n_layers, hidden_dim, output_dim, dropout):
        super().__init__()

        self.embedding = nn.Embedding(input_dim, emb_dim)
        self.positional_encoding = nn.Parameter(torch.zeros(1, 1000, emb_dim))
        self.transformer = nn.Transformer(emb_dim, n_heads, n_layers, n_layers, hidden_dim, dropout)
        self.fc = nn.Linear(emb_dim, output_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, src):
        src = src.permute(1, 0)  # (batch_size, seq_len) -> (seq_len, batch_size)
        embedded = self.dropout(self.embedding(src) + self.positional_encoding[:, :src.size(0), :])
        transformer_output = self.transformer(embedded, embedded)
        pooled = transformer_output.mean(dim=0)
        return self.fc(pooled)


INPUT_DIM = len(TEXT.vocab)
EMB_DIM = 256
N_HEADS = 8
N_LAYERS = 4
HIDDEN_DIM = 512
OUTPUT_DIM = 1
DROPOUT = 0.1

model = TransformerModel(INPUT_DIM, EMB_DIM, N_HEADS, N_LAYERS, HIDDEN_DIM, OUTPUT_DIM, DROPOUT).to(device)

optimizer = optim.Adam(model.parameters())
criterion = nn.BCEWithLogitsLoss().to(device)


def binary_accuracy(preds, y):
    rounded_preds = torch.round(torch.sigmoid(preds))
    correct = (rounded_preds == y).float()
    acc = correct.sum() / len(correct)
    return acc


def train(model, iterator, optimizer, criterion):
    epoch_loss = 0
    epoch_acc = 0
    model.train()

    for batch in iterator:
        optimizer.zero_grad()
        predictions = model(batch.text).squeeze(1)
        loss = criterion(predictions, batch.label)
        acc = binary_accuracy(predictions, batch.label)
        loss.backward()
        optimizer.step()

        epoch_loss += loss.item()
        epoch_acc += acc.item()

    return epoch_loss / len(iterator), epoch_acc / len(iterator)


def evaluate(model, iterator, criterion):
    epoch_loss = 0
    epoch_acc = 0
    model.eval()

    with torch.no_grad():
        for batch in iterator:
            predictions = model(batch.text).squeeze(1)
            loss = criterion(predictions, batch.label)
            acc = binary_accuracy(predictions, batch.label)

            epoch_loss += loss.item()
            epoch_acc += acc.item()

    return epoch_loss / len(iterator), epoch_acc / len(iterator)


N_EPOCHS = 5

for epoch in range(N_EPOCHS):
    train_loss, train_acc = train(model, train_iterator, optimizer, criterion)
    test_loss, test_acc = evaluate(model, test_iterator, criterion)

    print(f'Epoch: {epoch + 1}')
    print(f'\tTrain Loss: {train_loss:.3f} | Train Acc: {train_acc * 100:.2f}%')
    print(f'\tTest Loss: {test_loss:.3f} | Test Acc: {test_acc * 100:.2f}%')

nlp = spacy.load('en_core_web_sm')


def predict_sentiment(model, sentence):
    model.eval()
    tokenized = [tok.text for tok in nlp.tokenizer(sentence)]
    indexed = [TEXT.vocab.stoi[t] for t in tokenized]
    tensor = torch.LongTensor(indexed).to(device)
    tensor = tensor.unsqueeze(1)
    prediction = torch.sigmoid(model(tensor))
    return prediction.item()


sentence = "This film is terrible!"
print(f'Sentiment: {predict_sentiment(model, sentence):.3f}')

sentence = "This film is great!"
print(f'Sentiment: {predict_sentiment(model, sentence):.3f}')
